{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) MONAI Consortium\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\n",
    "from pathlib import Path\n",
    "import json\n",
    "from typing import Any, List, Dict\n",
    "import logging\n",
    "\n",
    "\n",
    "# Configure logging\n",
    "logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')\n",
    "\n",
    "\n",
    "def load_jsonl(path: Path) -> List[Dict[str, Any]]:\n",
    "    \"\"\"\n",
    "    Load data from a JSON Lines (JSONL) file.\n",
    "\n",
    "    Each line in the file should be a valid JSON object.\n",
    "\n",
    "    Args:\n",
    "        path (Path): The file path to the JSONL file.\n",
    "\n",
    "    Returns:\n",
    "        List[Dict[str, Any]]: A list of dictionaries parsed from the JSONL file.\n",
    "\n",
    "    Raises:\n",
    "        FileNotFoundError: If the specified file does not exist.\n",
    "        json.JSONDecodeError: If a line in the file is not valid JSON.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        with path.open('r', encoding='utf-8') as f:\n",
    "            return [json.loads(line) for line in f]\n",
    "    except FileNotFoundError:\n",
    "        raise FileNotFoundError(f\"The file {path} was not found.\")\n",
    "    except json.JSONDecodeError as e:\n",
    "        raise json.JSONDecodeError(f\"Invalid JSON in {path}: {e.msg}\", e.doc, e.pos)\n",
    "\n",
    "\n",
    "def load_json(path: Path) -> Any:\n",
    "    \"\"\"\n",
    "    Load data from a JSON file.\n",
    "\n",
    "    Args:\n",
    "        path (Path): The file path to the JSON file.\n",
    "\n",
    "    Returns:\n",
    "        Any: The data parsed from the JSON file.\n",
    "\n",
    "    Raises:\n",
    "        FileNotFoundError: If the specified file does not exist.\n",
    "        json.JSONDecodeError: If the file contains invalid JSON.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        with path.open('r', encoding='utf-8') as f:\n",
    "            return json.load(f)\n",
    "    except FileNotFoundError:\n",
    "        raise FileNotFoundError(f\"The file {path} was not found.\")\n",
    "    except json.JSONDecodeError as e:\n",
    "        raise json.JSONDecodeError(f\"Invalid JSON in {path}: {e.msg}\", e.doc, e.pos)\n",
    "\n",
    "\n",
    "def save_json(path: Path, data: Any) -> None:\n",
    "    \"\"\"\n",
    "    Save data to a JSON file with indentation for readability.\n",
    "\n",
    "    Args:\n",
    "        path (Path): The file path to save the JSON data.\n",
    "        data (Any): The data to be serialized and saved.\n",
    "\n",
    "    Raises:\n",
    "        IOError: If the file cannot be written.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        with path.open('w', encoding='utf-8') as f:\n",
    "            json.dump(data, f, indent=4)\n",
    "    except IOError as e:\n",
    "        raise IOError(f\"Failed to write to {path}: {e.strerror}\")\n",
    "\n",
    "\n",
    "def save_jsonl(path: Path, data: List[Dict[str, Any]]) -> None:\n",
    "    \"\"\"\n",
    "    Save a list of dictionaries to a JSON Lines (JSONL) file.\n",
    "\n",
    "    Each dictionary is serialized as a separate JSON object on its own line.\n",
    "\n",
    "    Args:\n",
    "        path (Path): The file path to save the JSONL data.\n",
    "        data (List[Dict[str, Any]]): A list of dictionaries to be serialized.\n",
    "\n",
    "    Raises:\n",
    "        IOError: If the file cannot be written.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        with path.open('w', encoding='utf-8') as f:\n",
    "            for line in data:\n",
    "                json.dump(line, f)\n",
    "                f.write('\\n')\n",
    "    except IOError as e:\n",
    "        raise IOError(f\"Failed to write to {path}: {e.strerror}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Process created json files for inference and evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running the vqa_json_creator.py script, we have json files containing multiple QA pairs for each image and consistent with the LLaVA-Med format. In this notebook, we will process these json files to create finalized json file for inference and evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add relative file path to the image ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_relative_path(input_file: Path, output_file: Path, path_table_file: Path) -> None:\n",
    "    \"\"\"\n",
    "    Add relative MIMIC image paths (including subdirectories) to JSON data based on a path table.\n",
    "\n",
    "    This function reads the input JSON file, replaces the 'image' field in each\n",
    "    entry with its corresponding relative path from the path table, and saves\n",
    "    the updated data to the output JSON file.\n",
    "\n",
    "    Args:\n",
    "        input_file (Path): Path to the input JSON file containing the data.\n",
    "        output_file (Path): Path to the output JSON file where updated data will be saved.\n",
    "        path_table_file (Path): Path to the JSON file mapping image IDs to their relative paths.\n",
    "\n",
    "    Raises:\n",
    "        KeyError: If an image ID in the input data is not found in the path table.\n",
    "        FileNotFoundError: If any of the specified files do not exist.\n",
    "        json.JSONDecodeError: If any of the JSON files contain invalid JSON.\n",
    "        IOError: If the output file cannot be written.\n",
    "    \"\"\"\n",
    "    data = load_json(input_file)\n",
    "    path_table = load_json(path_table_file)\n",
    "    \n",
    "    new_data = []\n",
    "    for entry in data:\n",
    "        image_id = entry.get('image')\n",
    "        if image_id is None:\n",
    "            raise KeyError(f\"Missing 'image' key in entry: {entry}\")\n",
    "        if image_id not in path_table:\n",
    "            raise KeyError(f\"Image ID '{image_id}' not found in path table.\")\n",
    "        new_entry = {\n",
    "            'image': path_table[image_id],\n",
    "            'id': entry.get('id'),\n",
    "            'conversations': list(entry.get('conversations', []))\n",
    "        }\n",
    "        new_data.append(new_entry)\n",
    "    \n",
    "    save_json(output_file, new_data)\n",
    "\n",
    "# Usage\n",
    "# input and output json files\n",
    "input_json = Path(\"input_json\") / \"llava_med_instruct_mimicvqa_test_expertmodel.json\"\n",
    "output_json = Path(\"output_json\") / \"llava_med_instruct_mimicvqa_test_expertmodel_path.json\"\n",
    "# json file mapping each image ID to its relative path\n",
    "path_table_json = Path(\"mimic_cxr_relpath.json\")\n",
    "add_relative_path(input_json, output_json, path_table_json)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split QA pairs into separate json files based on question type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_conversation(\n",
    "    old_list: List[Dict[str, Any]],\n",
    "    qtype: str,\n",
    "    qtable: Dict[str, str]\n",
    ") -> List[Dict[str, Any]]:\n",
    "    \"\"\"\n",
    "    Extract QA pairs from an old QA list that match a specific question type at the image level.\n",
    "    \n",
    "    Args:\n",
    "        old_list (List[Dict[str, Any]]): The original list of QA pairs.\n",
    "        qtype (str): The specific question type to filter by.\n",
    "        qtable (Dict[str, str]): A mapping of questions to their types.\n",
    "    \n",
    "    Returns:\n",
    "        List[Dict[str, Any]]: A new list containing only the QA pairs of the specified type.\n",
    "    \n",
    "    Raises:\n",
    "        ValueError: If the prefix format in the first QA pair is invalid.\n",
    "        KeyError: If a question is not found in the qtable.\n",
    "    \"\"\"\n",
    "    new_list = []\n",
    "    prefix = ''\n",
    "    suffix = '\\n<image>'\n",
    "\n",
    "    total_pairs = len(old_list) // 2\n",
    "    for i in range(total_pairs):\n",
    "        qa_index = i * 2\n",
    "        answer_index = qa_index + 1\n",
    "\n",
    "        # Process the first QA pair to determine the prefix\n",
    "        if i == 0:\n",
    "            text = old_list[qa_index].get('value', '')\n",
    "            test_split = text.split('\\n')\n",
    "            if len(test_split) == 2:\n",
    "                prefix = ''\n",
    "            elif len(test_split) == 3:\n",
    "                prefix = f\"{test_split[0]}\\n\"\n",
    "            else:\n",
    "                raise ValueError(f\"Invalid prefix format: {text}\")\n",
    "\n",
    "        # Extract and process the question\n",
    "        question_full = old_list[qa_index].get('value', '')\n",
    "        question_part = question_full.split('?')[0].split('\\n')[-1]\n",
    "        question_key = f\"{question_part}?\"\n",
    "\n",
    "        if question_key not in qtable:\n",
    "            raise KeyError(f\"Question '{question_key}' not found in question table.\")\n",
    "\n",
    "        if qtable[question_key] == qtype:\n",
    "            # Add the QA pair to the new list\n",
    "            if not new_list:\n",
    "                new_item = old_list[qa_index].copy()\n",
    "                new_item['value'] = f\"{prefix}{question_key}{suffix}\"\n",
    "                new_list.append(new_item)\n",
    "                new_list.append(old_list[answer_index].copy())\n",
    "            else:\n",
    "                new_list.append(old_list[qa_index].copy())\n",
    "                new_list.append(old_list[answer_index].copy())\n",
    "\n",
    "    return new_list\n",
    "\n",
    "\n",
    "def select_qa(\n",
    "    input_file: Path,\n",
    "    output_file: Path,\n",
    "    question_table_file: Path,\n",
    "    question_type: str\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Extract QA pairs from a JSON file that match a specific question type and save the filtered data.\n",
    "    \n",
    "    Args:\n",
    "        input_file (Path): Path to the input JSON file containing QA pairs.\n",
    "        output_file (Path): Path to the output JSON file to save filtered QA pairs.\n",
    "        question_table_file (Path): Path to the JSON file mapping each question to its type.\n",
    "        question_type (str): The specific question type to filter by.\n",
    "    \n",
    "    Raises:\n",
    "        FileNotFoundError: If any of the specified files do not exist.\n",
    "        json.JSONDecodeError: If any of the JSON files contain invalid JSON.\n",
    "        KeyError: If a question is not found in the question table.\n",
    "        IOError: If the output file cannot be written.\n",
    "    \"\"\"\n",
    "    data = load_json(input_file)\n",
    "    qtable = load_json(question_table_file)\n",
    "    new_data = []\n",
    "\n",
    "    for entry in data:\n",
    "        conversations = entry.get('conversations', [])\n",
    "        filtered_conversations = select_conversation(conversations, question_type, qtable)\n",
    "        if filtered_conversations:\n",
    "            new_entry = {\n",
    "                'image': entry.get('image'),\n",
    "                'id': entry.get('id'),\n",
    "                'conversations': filtered_conversations\n",
    "            }\n",
    "            new_data.append(new_entry)\n",
    "\n",
    "    save_json(output_file, new_data)\n",
    "\n",
    "# Usage\n",
    "question_types = ['abnormality', 'level', 'location', 'presence', 'type', 'view']\n",
    "question_table_file = Path('mimic_vqa_qtype.json')\n",
    "root = Path(\"output_json\")\n",
    "\n",
    "for qtype in question_types:\n",
    "    output_dir = root / qtype\n",
    "    if not output_dir.exists():\n",
    "        output_dir.mkdir(parents=True)\n",
    "    \n",
    "    # with expert model\n",
    "    input_file = root / 'llava_med_instruct_mimicvqa_test_expertmodel_path.json'\n",
    "    output_file = output_dir / 'llava_med_instruct_mimicvqa_test_expert.json'\n",
    "    select_qa(input_file, output_file, question_table_file, qtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These JSON files can now be used for first-round inference. However, during training, expert information is only included in the first question of each QA conversation round. Therefore, to generate the inference correctly, the following changes need to be made:\n",
    "\n",
    "* If a question q is not the first question in a QA conversation, a second-round inference is required.\n",
    "* In this second-round inference, the first question in that QA conversation, along with its predicted answer from the first-round inference and the expert information, should be inserted before the question q."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dict_id_predictions(input_file: Path) -> Dict[str, Any]:\n",
    "    \"\"\"\n",
    "    Convert first-round inference answers to a dictionary with combined keys.\n",
    "\n",
    "    Each key is a combination of 'question_id' and 'prompt', and the value is the corresponding 'text'.\n",
    "\n",
    "    Args:\n",
    "        input_file (Path): Path to the input JSONL file containing inference answers.\n",
    "\n",
    "    Returns:\n",
    "        Dict[str, Any]: A dictionary mapping combined keys to their prediction texts.\n",
    "\n",
    "    Raises:\n",
    "        KeyError: If expected keys are missing in any data entry.\n",
    "    \"\"\"\n",
    "    data = load_jsonl(input_file)\n",
    "    res = {}\n",
    "    for entry in data:\n",
    "        try:\n",
    "            question_id = entry['question_id']\n",
    "            prompt = entry['prompt']\n",
    "            text = entry['text']\n",
    "            key = f\"{question_id}_{prompt}\"\n",
    "            res[key] = text\n",
    "        except KeyError as e:\n",
    "            logging.warning(f\"Missing key {e} in entry: {entry}\")\n",
    "            continue\n",
    "    return res\n",
    "\n",
    "\n",
    "# Usage\n",
    "# Configure logging\n",
    "logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')\n",
    "question_types = ['abnormality', 'presence', 'view', 'location', 'level', 'type']\n",
    "predictions_root = Path(\"../predicted_vqa_json/ckpt_llava_med_all_mimic_expert_1run\")\n",
    "predictions = {}\n",
    "for qtype in question_types:\n",
    "    input_file = predictions_root / qtype / 'llava_med_expert_all_mimic_expert_run1.jsonl'\n",
    "    predictions.update(create_dict_id_predictions(input_file))\n",
    "\n",
    "save_json(Path('predictions.json'), predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_expert_prompt(\n",
    "    image_id: str,\n",
    "    old_list: List[Dict[str, Any]],\n",
    "    qtype: str,\n",
    "    qtable: Dict[str, str],\n",
    "    predictions: Dict[str, Any]\n",
    ") -> List[Dict[str, Any]]:\n",
    "    \"\"\"\n",
    "    Insert first-round inference (predictions) and expert information into each QA pair.\n",
    "    \n",
    "    This function processes a list of QA pairs, updates the first QA pair with prediction\n",
    "    and expert information based on the question type, and filters QA pairs that match\n",
    "    the specified question type.\n",
    "    \n",
    "    Args:\n",
    "        image_id (str): The unique identifier for the image.\n",
    "        old_list (List[Dict[str, Any]]): The original list of QA pairs.\n",
    "        qtype (str): The specific question type to filter by.\n",
    "        qtable (Dict[str, str]): A mapping of questions to their types.\n",
    "        predictions (Dict[str, Any]): A dictionary containing prediction texts keyed by\n",
    "                                      a combination of image_id and question prompt.\n",
    "    \n",
    "    Returns:\n",
    "        List[Dict[str, Any]]: A new list of QA pairs with inserted prediction and expert information.\n",
    "    \n",
    "    Raises:\n",
    "        ValueError: If the first question prompt has an invalid format.\n",
    "        KeyError: If a question is not found in the qtable or predictions.\n",
    "    \"\"\"\n",
    "    new_list = []\n",
    "    prefix = ''\n",
    "    suffix = ''\n",
    "    total_pairs = len(old_list) // 2\n",
    "    \n",
    "    for i in range(total_pairs):\n",
    "        qa_index = i * 2\n",
    "        answer_index = qa_index + 1\n",
    "        \n",
    "        if i == 0:\n",
    "            qa_pair = old_list[qa_index]\n",
    "            text = qa_pair.get('value', '')\n",
    "            text_split = text.split('\\n')\n",
    "            \n",
    "            try:\n",
    "                # Construct the key to retrieve the prediction\n",
    "                pred_answer = predictions[f\"{image_id}_{text}\"]\n",
    "            except KeyError:\n",
    "                logging.warning(f\"Prediction not found for key: {image_id}_{text}\")\n",
    "                continue\n",
    "            \n",
    "            if len(text_split) == 2:\n",
    "                expert = ''\n",
    "                question = text_split[0]\n",
    "            elif len(text_split) == 3:\n",
    "                expert = f\"{text_split[0]}\\n\"\n",
    "                question = text_split[1]\n",
    "            else:\n",
    "                raise ValueError(f\"Invalid first question prompt format: {text}\")\n",
    "            \n",
    "            prefix = f\"{expert}{question}\\n<patch_token_holder>\\n###{pred_answer}\\n###Human: \"\n",
    "            \n",
    "            # Check if the question type matches\n",
    "            question_key = f\"{question}\"\n",
    "            if qtable.get(question_key) == qtype:\n",
    "                updated_qa = qa_pair.copy()\n",
    "                updated_qa['value'] = f\"{expert}{question}\\n<patch_token_holder>\"\n",
    "                new_list.append(updated_qa)\n",
    "                new_list.append(old_list[answer_index].copy())\n",
    "        else:\n",
    "            qa_pair = old_list[qa_index]\n",
    "            question = qa_pair.get('value', '')\n",
    "            \n",
    "            if qtable.get(question) == qtype:\n",
    "                if not new_list:\n",
    "                    updated_qa = qa_pair.copy()\n",
    "                    updated_qa['value'] = f\"{prefix}{question}{suffix}\"\n",
    "                    new_list.append(updated_qa)\n",
    "                    new_list.append(old_list[answer_index].copy())\n",
    "                else:\n",
    "                    new_list.append(qa_pair.copy())\n",
    "                    new_list.append(old_list[answer_index].copy())\n",
    "                    \n",
    "    return new_list\n",
    "\n",
    "\n",
    "def add_expert_prompt_json(\n",
    "    input_file: Path,\n",
    "    output_file: Path,\n",
    "    question_table_file: Path,\n",
    "    question_type: str,\n",
    "    predictions_file: Path\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Apply expert prompts to all QA pairs in a JSON file based on question type and predictions.\n",
    "    \n",
    "    This function processes each entry in the input JSON file, updates the conversations\n",
    "    with expert prompts where applicable, and saves the filtered and updated data to the\n",
    "    output JSON file.\n",
    "    \n",
    "    Args:\n",
    "        input_file (Path): Path to the input JSON file containing QA pairs.\n",
    "        output_file (Path): Path to the output JSON file to save updated QA pairs.\n",
    "        question_table_file (Path): Path to the JSON file mapping questions to their types.\n",
    "        question_type (str): The specific question type to filter by.\n",
    "        predictions_file (Path): Path to the JSON file containing prediction texts.\n",
    "    \n",
    "    Raises:\n",
    "        FileNotFoundError: If any of the specified files do not exist.\n",
    "        json.JSONDecodeError: If any of the JSON files contain invalid JSON.\n",
    "        KeyError: If a question is not found in the question table or predictions.\n",
    "        IOError: If the output file cannot be written.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        data = load_json(input_file)\n",
    "        qtable = load_json(question_table_file)\n",
    "        predictions = load_json(predictions_file)\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Failed to load JSON files: {e}\")\n",
    "        raise\n",
    "    \n",
    "    new_data = []\n",
    "    for entry in data:\n",
    "        image_id = entry.get('id')\n",
    "        conversations = entry.get('conversations', [])\n",
    "        \n",
    "        if not image_id:\n",
    "            logging.warning(f\"Missing 'id' in entry: {entry}\")\n",
    "            continue\n",
    "        \n",
    "        try:\n",
    "            updated_conversations = add_expert_prompt(\n",
    "                image_id=image_id,\n",
    "                old_list=conversations,\n",
    "                qtype=question_type,\n",
    "                qtable=qtable,\n",
    "                predictions=predictions\n",
    "            )\n",
    "            if updated_conversations:\n",
    "                new_entry = {\n",
    "                    'image': entry.get('image'),\n",
    "                    'id': image_id,\n",
    "                    'conversations': updated_conversations\n",
    "                }\n",
    "                new_data.append(new_entry)\n",
    "        except Exception as e:\n",
    "            logging.error(f\"Failed to process entry with id {image_id}: {e}\")\n",
    "            continue\n",
    "    \n",
    "    try:\n",
    "        save_json(output_file, new_data)\n",
    "        logging.info(f\"Successfully saved updated data to {output_file}.\")\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Failed to save output file {output_file}: {e}\")\n",
    "        raise\n",
    "\n",
    "# Usage\n",
    "# Configure logging\n",
    "logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')\n",
    "question_types = ['abnormality', 'presence', 'view', 'location', 'level', 'type']\n",
    "question_table_file = Path('mimic_vqa_qtype.json')\n",
    "predictions_file = Path('predictions.json')\n",
    "input_file = Path('output_json') / 'llava_med_instruct_mimicvqa_test_expertmodel_path.json'\n",
    "for qtype in question_types:\n",
    "    output_dir = Path('output_json') / qtype\n",
    "    if not output_dir.exists():\n",
    "        output_dir.mkdir(parents=True)\n",
    "    \n",
    "    # with expert model\n",
    "    output_file = output_dir / 'llava_med_instruct_mimicvqa_test_expert_2run.json'\n",
    "    add_expert_prompt_json(input_file, output_file, question_table_file, qtype, predictions_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add question and answer type to testing json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_qa_type(\n",
    "    input_file: Path,\n",
    "    output_file: Path,\n",
    "    question_type: str\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Insert question type and determine answer type for each QA pair in a JSON file.\n",
    "    \n",
    "    This function processes each entry in the input JSON file, analyzes the answers to\n",
    "    categorize them as 'closed' or 'open', and appends the question type and answer type\n",
    "    to each entry. The updated data is then saved to the output JSON file.\n",
    "    \n",
    "    Args:\n",
    "        input_file (Path): Path to the input JSON file containing QA pairs.\n",
    "        output_file (Path): Path to the output JSON file to save updated QA pairs.\n",
    "        question_type (str): The specific question type to assign to each entry.\n",
    "    \n",
    "    Raises:\n",
    "        FileNotFoundError: If the input file does not exist.\n",
    "        json.JSONDecodeError: If the input file contains invalid JSON.\n",
    "        IOError: If the output file cannot be written.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        data = load_json(input_file)\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Failed to load input file {input_file}: {e}\")\n",
    "        raise\n",
    "\n",
    "    new_data = []\n",
    "    \n",
    "    for entry in data:\n",
    "        try:\n",
    "            conversations = entry.get('conversations', [])\n",
    "            if len(conversations) < 2:\n",
    "                logging.warning(f\"Entry with id {entry.get('id')} has insufficient conversations.\")\n",
    "                continue\n",
    "\n",
    "            # Extract the question text\n",
    "            question_full = conversations[0].get('value', '')\n",
    "            question_split = question_full.split('?')[0].split('\\n')\n",
    "            question = question_split[-1] if question_split else ''\n",
    "\n",
    "            # Extract and process the answer\n",
    "            answer = conversations[1].get('value', '').lower().strip()\n",
    "            \n",
    "            if answer in {'yes', 'no'}:\n",
    "                answer_type = 'closed'\n",
    "            elif 'yes' in answer or 'no' in answer:\n",
    "                answer_type = 'closed'\n",
    "                logging.warning(f\"Ambiguous answer for question '{question}': '{answer}'\")\n",
    "            else:\n",
    "                answer_type = 'open'\n",
    "            \n",
    "            # Append the updated entry\n",
    "            new_entry = {\n",
    "                'image': entry.get('image'),\n",
    "                'id': entry.get('id'),\n",
    "                'conversations': conversations,\n",
    "                'question_type': question_type,\n",
    "                'answer_type': answer_type\n",
    "            }\n",
    "            new_data.append(new_entry)\n",
    "        \n",
    "        except Exception as e:\n",
    "            logging.error(f\"Error processing entry with id {entry.get('id')}: {e}\")\n",
    "            continue\n",
    "    \n",
    "    try:\n",
    "        save_json(output_file, new_data)\n",
    "        logging.info(f\"Successfully saved updated data to {output_file}.\")\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Failed to save output file {output_file}: {e}\")\n",
    "        raise\n",
    "\n",
    "\n",
    "# Usage\n",
    "# Configure logging\n",
    "logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')\n",
    "    \n",
    "question_types = ['abnormality', 'level', 'location', 'presence', 'type', 'view']\n",
    "root = Path('output_json')\n",
    "\n",
    "for qtype in question_types:\n",
    "    output_dir = root / qtype\n",
    "    if not output_dir.exists():\n",
    "        output_dir.mkdir(parents=True)\n",
    "    \n",
    "    # with expert model\n",
    "    input_file = output_dir / 'llava_med_instruct_mimicvqa_test_expert.json'\n",
    "    output_file = output_dir / 'llava_med_instruct_mimicvqa_test_expert_type.json'\n",
    "    add_qa_type(input_file, output_file, qtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
